{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CatTSunami Tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fairchem.applications.cattsunami.core import Reaction\n",
    "from fairchem.data.oc.core import Slab, Adsorbate, Bulk, AdsorbateSlabConfig\n",
    "from fairchem.core.common.relaxation.ase_utils import OCPCalculator\n",
    "from ase.optimize import BFGS\n",
    "from x3dase.visualize import view_x3d_n\n",
    "from ase.io import read\n",
    "from x3dase.x3d import X3D\n",
    "from fairchem.applications.cattsunami.databases import DISSOCIATION_REACTION_DB_PATH\n",
    "from fairchem.data.oc.databases.pkls import ADSORBATE_PKL_PATH, BULK_PKL_PATH\n",
    "from fairchem.core.models.model_registry import model_name_to_local_file\n",
    "import matplotlib.pyplot as plt\n",
    "from fairchem.applications.cattsunami.core.autoframe import AutoFrameDissociation\n",
    "from fairchem.applications.cattsunami.core import OCPNEB\n",
    "from ase.io import read\n",
    "\n",
    "#Optional\n",
    "from IPython.display import Image\n",
    "from x3dase.x3d import X3D\n",
    "\n",
    "#Set random seed\n",
    "import numpy as np\n",
    "np.random.seed(22)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Do enumerations in an AdsorbML style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the reaction class for the reaction of interest\n",
    "reaction = Reaction(reaction_str_from_db=\"*CH -> *C + *H\",\n",
    "                    reaction_db_path=DISSOCIATION_REACTION_DB_PATH,\n",
    "                    adsorbate_db_path = ADSORBATE_PKL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate our adsorbate class for the reactant and product\n",
    "reactant = Adsorbate(adsorbate_id_from_db=reaction.reactant1_idx, adsorbate_db_path=ADSORBATE_PKL_PATH)\n",
    "product1 = Adsorbate(adsorbate_id_from_db=reaction.product1_idx, adsorbate_db_path=ADSORBATE_PKL_PATH)\n",
    "product2 = Adsorbate(adsorbate_id_from_db=reaction.product2_idx, adsorbate_db_path=ADSORBATE_PKL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grab the bulk and cut the slab we are interested in\n",
    "bulk = Bulk(bulk_src_id_from_db=\"mp-33\", bulk_db_path=BULK_PKL_PATH)\n",
    "slab = Slab.from_bulk_get_specific_millers(bulk = bulk, specific_millers=(0,0,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform site enumeration\n",
    "# For AdsorbML num_sites = 100, but we use 5 here for brevity. This should be increased for practical use.\n",
    "reactant_configs = AdsorbateSlabConfig(slab = slab[0], adsorbate = reactant,\n",
    "                                       mode=\"random_site_heuristic_placement\",\n",
    "                                       num_sites = 5).atoms_list\n",
    "product1_configs = AdsorbateSlabConfig(slab = slab[0], adsorbate = product1,\n",
    "                                      mode=\"random_site_heuristic_placement\",\n",
    "                                      num_sites = 5).atoms_list\n",
    "product2_configs = AdsorbateSlabConfig(slab = slab[0], adsorbate = product2,\n",
    "                                      mode=\"random_site_heuristic_placement\",\n",
    "                                      num_sites = 5).atoms_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the calculator\n",
    "# NOTE: If you have a GPU, use cpu = False\n",
    "# NOTE: Change the checkpoint path to locally downloaded files as needed\n",
    "checkpoint_path = model_name_to_local_file('EquiformerV2-31M-S2EF-OC20-All+MD', local_cache='/tmp/ocp_checkpoints/')\n",
    "cpu = True\n",
    "calc = OCPCalculator(checkpoint_path = checkpoint_path, cpu = cpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Relax the reactant systems\n",
    "reactant_energies = []\n",
    "for config in reactant_configs:\n",
    "    config.calc = calc\n",
    "    opt = BFGS(config)\n",
    "    opt.run(fmax = 0.05, steps=200)\n",
    "    reactant_energies.append(config.get_potential_energy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Relax the product systems\n",
    "product1_energies = []\n",
    "for config in product1_configs:\n",
    "    config.calc = calc\n",
    "    opt = BFGS(config)\n",
    "    opt.run(fmax = 0.05, steps=200)\n",
    "    product1_energies.append(config.get_potential_energy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "product2_energies = []\n",
    "for config in product2_configs:\n",
    "    config.calc = calc\n",
    "    opt = BFGS(config)\n",
    "    opt.run(fmax = 0.05, steps=200)\n",
    "    product2_energies.append(config.get_potential_energy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enumerate NEBs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=\"dissociation_scheme.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "af = AutoFrameDissociation(\n",
    "            reaction = reaction,\n",
    "            reactant_system = reactant_configs[reactant_energies.index(min(reactant_energies))],\n",
    "            product1_systems = product1_configs,\n",
    "            product1_energies = product1_energies,\n",
    "            product2_systems = product2_configs,\n",
    "            product2_energies = product2_energies,\n",
    "            r_product1_max=2, #r1 in the above fig\n",
    "            r_product2_max=3, #r3 in the above fig\n",
    "            r_product2_min=1, #r2 in the above fig\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nframes = 10\n",
    "frame_sets, mapping_idxs = af.get_neb_frames(calc,\n",
    "                               n_frames = nframes,\n",
    "                               n_pdt1_sites=4, # = 5 in the above fig (step 1)\n",
    "                               n_pdt2_sites = 4, # = 5 in the above fig (step 2)\n",
    "                              )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run NEBs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## This will run all NEBs enumerated - to just run one, run the code cell below.\n",
    "# On GPU, each NEB takes an average of ~1 minute so this could take around a half hour on GPU\n",
    "# But much longer on CPU\n",
    "# Remember that not all NEBs will converge -- the k, nframes would be adjusted to achieve convergence\n",
    "\n",
    "# fmax = 0.05 # [eV / ang**2]\n",
    "# delta_fmax_climb = 0.4\n",
    "# converged_idxs = []\n",
    "\n",
    "# for idx, frame_set in enumerate(frame_sets):\n",
    "#     neb = OCPNEB(\n",
    "#         frame_set,\n",
    "#         checkpoint_path=checkpoint_path,\n",
    "#         k=1,\n",
    "#         batch_size=8,\n",
    "#         cpu = cpu,\n",
    "#     )\n",
    "#     optimizer = BFGS(\n",
    "#         neb,\n",
    "#         trajectory=f\"ch_dissoc_on_Ru_{idx}.traj\",\n",
    "#     )\n",
    "#     conv = optimizer.run(fmax=fmax + delta_fmax_climb, steps=200)\n",
    "#     if conv:\n",
    "#         neb.climb = True\n",
    "#         conv = optimizer.run(fmax=fmax, steps=300)\n",
    "#         if conv:\n",
    "#             converged_idxs.append(idx)\n",
    "            \n",
    "# print(converged_idxs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you run the above cell -- dont run this one\n",
    "fmax = 0.05 # [eV / ang**2]\n",
    "delta_fmax_climb = 0.4\n",
    "neb = OCPNEB(\n",
    "    frame_sets[0],\n",
    "    checkpoint_path=checkpoint_path,\n",
    "    k=1,\n",
    "    batch_size=8,\n",
    "    cpu = cpu,\n",
    ")\n",
    "optimizer = BFGS(\n",
    "    neb,\n",
    "    trajectory=f\"ch_dissoc_on_Ru_0.traj\",\n",
    ")\n",
    "conv = optimizer.run(fmax=fmax + delta_fmax_climb, steps=200)\n",
    "if conv:\n",
    "    neb.climb = True\n",
    "    conv = optimizer.run(fmax=fmax, steps=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_neb = read(f\"ch_dissoc_on_Ru_{converged_idxs[0]}.traj\", \":\")[-1*nframes:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "es  = []\n",
    "for frame in optimized_neb:\n",
    "    frame.set_calculator(calc)\n",
    "    es.append(frame.get_potential_energy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the reaction coordinate\n",
    "\n",
    "es = [e - es[0] for e in es]\n",
    "plt.plot(es)\n",
    "plt.xlabel(\"frame number\")\n",
    "plt.ylabel(\"relative energy [eV]\")\n",
    "plt.title(f\"CH dissociation on Ru(0001), Ea = {max(es):1.2f} eV\")\n",
    "plt.savefig(\"CH_dissoc_on_Ru_0001.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make an interative html file of the optimized neb trajectory\n",
    "x3d = X3D(optimized_neb)\n",
    "x3d.write(\"optimized_neb_ch_disoc_on_Ru0001.html\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
